from core.layer import Layer
from implement.layers.conv.bottleneckA import BottleneckA
from implement.layers.conv.bottleneckB import BottleneckB

'''
BuildingBlock 类是一个构建深度神经网络的基本模块。它是由多个堆叠在一起的残差块（Residual Block）组成，每个残差块包含多个子层。
在深度神经网络中，通过堆叠多个这样的基本模块，可以构建出更复杂的神经网络结构。
这个特定的 BuildingBlock 类由两种不同的残差块组成，分别是 BottleneckA 和 BottleneckB。
这两种块具有不同的结构，通过组合它们，可以构建出层次更深、表达能力更强大的神经网络。
总的来说，BuildingBlock 类的作用是提供一个可重复使用的、模块化的网络结构，使得神经网络的设计和构建更加方便和灵活。
通过调整其中的参数，可以很容易地搭建不同深度和复杂度的神经网络。
'''
class BuildingBlock(Layer):

    def __init__(self, n_layers=None, in_channels=None, mid_channels=None,
                 out_channels=None, stride=None, downsample_fb=None):
        """
        初始化 BuildingBlock 类。

        Parameters:
            n_layers (int): Block 中的层数。
            in_channels (int): 输入通道数。
            mid_channels (int): 中间层通道数。
            out_channels (int): 输出通道数。
            stride (int): 步幅。
            downsample_fb (bool): 是否进行下采样。
        """
        super().__init__()

        # 创建 BottleneckA 层
        self.a = BottleneckA(in_channels, mid_channels, out_channels, stride,
                             downsample_fb)
        self._forward = ['a']

        # 创建多个 BottleneckB 层
        for i in range(n_layers - 1):
            name = 'b{}'.format(i + 1)
            bottleneck = BottleneckB(out_channels, mid_channels)
            setattr(self, name, bottleneck)
            self._forward.append(name)

    def forward(self, x):
        """
        前向传播函数，将输入数据经过 BuildingBlock 处理。

        Parameters:
            x: 输入数据。

        Returns:
            y: 处理后的数据。
        """
        for name in self._forward:
            l = getattr(self, name)
            x = l(x)
        return x

